import math
import unittest
from electrostaticvacuum import Vector, PointCharge, Field


class ElectrostaticFieldTest(unittest.TestCase):
    def assertPointChargeListAlmostEqual(self, lst1: list[PointCharge], lst2: list[PointCharge], delta: float = 1E-6):
        self.assertEqual(len(lst1), len(lst2))
        for left, right in zip(lst1, lst2):
            self.assertAlmostEqual(left.x, right.x, delta=delta)
            self.assertAlmostEqual(left.y, right.y, delta=delta)
            self.assertAlmostEqual(left.charge, right.charge, delta=delta)

    def test_of_x_poles(self):
        field = Field.of_x_poles(0, k=1123)
        self.assertPointChargeListAlmostEqual(
            [],
            list(field.get_charges())
        )
        self.assertAlmostEqual(1123, field.k)

        field = Field.of_x_poles(1)
        self.assertPointChargeListAlmostEqual(
            [PointCharge(Vector(1, 0), 1)],
            list(field.get_charges())
        )
        self.assertAlmostEqual(8.99E9, field.k)

        field = Field.of_x_poles(4, r=21.213, charge=11.21, start_angle=math.pi / 2, k=11123)
        self.assertPointChargeListAlmostEqual(
            [
                PointCharge(Vector(0, 21.213), 11.21),
                PointCharge(Vector(-21.213, 0), -11.21),
                PointCharge(Vector(0, -21.213), 11.21),
                PointCharge(Vector(21.213, 0), -11.21)
            ],
            list(field.get_charges())
        )
        self.assertAlmostEqual(11123, field.k)

    def test_from_charge_map(self):
        charge_map = [
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, -123, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
        ]
        field = Field.from_charge_map(charge_map, Vector(-0.3, -0.3), 0.1, 0.1, k=1234)
        self.assertPointChargeListAlmostEqual(
            [
                PointCharge(Vector(-0.1, -0.1), 1),
                PointCharge(Vector(0.1, 0.1), -123)
            ],
            list(field.get_charges())
        )
        self.assertAlmostEqual(1234, field.k)


if __name__ == '__main__':
    unittest.main()
